
#include "utils.h"

extern "C" aclrtRunMode runNode()
{
    aclrtRunMode mode;
    aclrtGetRunMode(&mode);
    return mode;
}

aclEnv::aclEnv(int32_t deviceId) : device_id(deviceId)
{
    // 设置Device
    auto ret = aclrtSetDevice(device_id);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret));
    // 创建Context
    ret = aclrtCreateContext(&context, device_id);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret));
    // 设置Context
    ret = aclrtSetCurrentContext(context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret));
    // 设置Stream
    ret = aclrtCreateStream(&stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret));
    // 执行初始化
    ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret));
}

aclEnv::~aclEnv()
{
    aclrtDestroyStream(stream);
    aclrtDestroyContext(context);
    aclrtResetDevice(device_id);
    aclFinalize();
}

int aclEnv::syncStream()
{
    return aclrtSynchronizeStream(stream);
}

NDTensor::NDTensor(void *shape, uint64_t shape_size, int dataType, void *tensorData)
{
    hostAddr = tensorData;
    auto shape_int64 = (int64_t *)shape;
    auto acl_data_type = (aclDataType)dataType;
    
    // 获取size
    auto desc = aclCreateTensorDesc(acl_data_type, shape_size, shape_int64, aclFormat::ACL_FORMAT_ND);
    size = aclGetTensorDescSize(desc);
    aclDestroyTensorDesc(desc);

    // 分配设备内存
    auto ret = aclrtMalloc(&deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret));

    // 拷贝数据
    aclrtMemcpy(deviceAddr, size, hostAddr, size, aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE);

    // 创建tensor
    tensorPtr = aclCreateTensor(
        shape_int64,              // view shape的维度值。
        shape_size,               // view shape的维度个数
        acl_data_type,            // 数据类型
        nullptr,                  // 各维度的内存跨度数据
        0,                        // tensor首元素相对于storage的偏移。
        aclFormat::ACL_FORMAT_ND, // tensor内存排布格式
        shape_int64,              // tensor的存储shape的维度值。
        shape_size,               // tensor的存储shape的维度数。
        deviceAddr                // tensor在Device侧的存储地址。
    );

    // 销毁描述
}

NDTensor::~NDTensor()
{
    if (deviceAddr != nullptr)
    {
        aclrtFree(deviceAddr);
        deviceAddr = nullptr;
    }
}

void NDTensor::to_cpu()
{
    aclrtMemcpy(hostAddr, size, deviceAddr, size, aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST);
}

extern "C"
{
    void *creatEnv(int32_t device_id)
    {
        return new aclEnv(device_id);
    }

    void destroyEnv(void *p)
    {
        delete ((aclEnv *)p);
    }

    int envSyncStream(void *p)
    {
        return ((aclEnv *)p)->syncStream();
    }

    void *createNDTensor(void *shape, uint64_t shape_size, int dataType, void *tensorData)
    {
        return new NDTensor(shape, shape_size, dataType, tensorData);
    }

    void destroyNDTensor(void *p)
    {
        delete ((NDTensor *)p);
    }

    void ndTensor2Cpu(void *p)
    {
        ((NDTensor *)p)->to_cpu();
    }
}
